#!/usr/bin/python3
# -*- coding:utf-8 -*-
import os
import sys
import logging
import numpy as np

# ====================================
LOSS      = 2e-3      # 容差
MINIMUM   = 1e-10     # 防止除零
REDUCTION = 0         # 0 = 全部对比，1 = 只比第一个值
BATCH     = 8
FEATURES  = 2048
# ====================================

logging.basicConfig(
    level=logging.DEBUG,
    format='%(asctime)s - %(levelname)s - %(message)s'
)

def load_complex_data(path):
    """从 float32 文件中读取数据，并转换成 complex64"""
    data = np.fromfile(path, dtype=np.float32)
    if data.size != BATCH * FEATURES * 2:
        logging.error(f"文件 {path} 尺寸不符合预期: {data.size} vs {BATCH*FEATURES*2}")
        sys.exit(1)
    data_complex = data.view(np.complex64).reshape(BATCH, FEATURES)
    return data_complex

def format_complex(c):
    """格式化复数为 a+bj"""
    return f"{c.real:.6f} + {c.imag:.6f}j"

def verify_first_only(real_path, golden_path):
    real = load_complex_data(real_path)[0,0]
    golden = load_complex_data(golden_path)[0,0]
    abs_err = abs(real - golden)
    rel_err = abs_err / (abs(golden) + MINIMUM)
    logging.info(f"First value  real : {format_complex(real)}")
    logging.info(f"First value golden: {format_complex(golden)}")
    if abs_err <= LOSS and rel_err <= LOSS:
        logging.info("FIRST VALUE: PASS — within tolerance.")
        sys.exit(0)
    else:
        logging.error("FIRST VALUE: FAIL — out of tolerance.")
        logging.error(f"abs_err = {abs_err}, rel_err = {rel_err}")
        sys.exit(1)

def verify_all(real_path, golden_path):
    real = load_complex_data(real_path)
    golden = load_complex_data(golden_path)
    abs_err = np.abs(real - golden)
    rel_err = abs_err / (np.abs(golden) + MINIMUM)
    pass_mask = (abs_err <= LOSS) & (rel_err <= LOSS)
    if np.all(pass_mask):
        logging.info("ALL ELEMENTS: PASS — within tolerance.")
        sys.exit(0)
    else:
        bad_idx = np.argwhere(~pass_mask)
        logging.error(f"ELEMENTS FAIL: {len(bad_idx)} / {real.size}")
        for i, j in bad_idx[:10]:
            logging.error(f"Idx ({i},{j}): real={format_complex(real[i,j])}, "
                          f"golden={format_complex(golden[i,j])}, "
                          f"abs_err={abs_err[i,j]}, rel_err={rel_err[i,j]}")
        sys.exit(1)

def main():
    if len(sys.argv) != 3:
        logging.error("Usage: python verify.py <real_result.bin> <golden.bin>")
        sys.exit(1)

    logging.info(f"REDUCTION = {REDUCTION}  (0=all, 1=first)")
    if REDUCTION == 0:
        verify_all(sys.argv[1], sys.argv[2])
    else:
        verify_first_only(sys.argv[1], sys.argv[2])

if __name__ == '__main__':
    main()
